package com.pig4cloud.pigx.knowledge.support.rule;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.date.DateUtil;
import cn.hutool.core.util.ReflectUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.json.JSONUtil;
import com.pig4cloud.pigx.common.core.util.R;
import com.pig4cloud.pigx.knowledge.dto.AiMessageResultDTO;
import com.pig4cloud.pigx.knowledge.dto.BaseAiRequest;
import com.pig4cloud.pigx.knowledge.dto.ChatMessageDTO;
import com.pig4cloud.pigx.knowledge.service.AiAssistantService;
import com.pig4cloud.pigx.knowledge.support.function.FunctionCalling;
import com.pig4cloud.pigx.knowledge.support.provider.ChatMemoryAdvisorProvider;
import com.pig4cloud.pigx.knowledge.support.provider.ModelProvider;
import com.pig4cloud.pigx.knowledge.support.util.ChatMessageContextHolder;
import com.pig4cloud.pigx.knowledge.support.util.PromptBuilder;
import com.pig4cloud.pigx.knowledge.support.util.ToolSpecificationsUtils;
import dev.ai4j.openai4j.chat.ResponseFormatType;
import dev.langchain4j.agent.tool.ToolExecutionRequest;
import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.memory.ChatMemory;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.chat.request.ChatRequest;
import dev.langchain4j.model.chat.request.ResponseFormat;
import dev.langchain4j.model.chat.request.json.JsonObjectSchema;
import dev.langchain4j.model.chat.request.json.JsonSchema;
import dev.langchain4j.model.chat.response.ChatResponse;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.tuple.Triple;
import org.springframework.stereotype.Component;
import reactor.core.publisher.Flux;

import java.lang.reflect.Field;
import java.util.*;

import static dev.langchain4j.data.message.SystemMessage.systemMessage;
import static dev.langchain4j.data.message.UserMessage.userMessage;
import static dev.langchain4j.model.chat.request.ResponseFormatType.JSON;

/**
 * 基于JSON 结构化交互
 *
 * @author lengleng
 * @date 2024/3/26
 */
@Slf4j
@Component("jsonChat")
@RequiredArgsConstructor
public class JsonChatRule implements ChatRule {

    private final ChatMemoryAdvisorProvider chatMemoryAdvisorProvider;

    private final List<FunctionCalling> functionCallingList;

    private final ModelProvider modelProvider;

    /**
     * 处理聊天信息
     *
     * @param chatMessageDTO 聊天上文
     * @return flux stream
     */
    @Override
    public Flux<AiMessageResultDTO> process(ChatMessageDTO chatMessageDTO) {

        Triple<ChatLanguageModel, AiAssistantService, String> jsonAssistantTriple = modelProvider
                .getAiJSONAssistant(chatMessageDTO.getModelName());

        // 判断是否包含函数名称
        Optional<FunctionCalling> functionCallingOptional = functionCallingList.stream()
                .filter(functionCalling -> StrUtil.contains(chatMessageDTO.getExtDetails().getFuncName(),
                        functionCalling.functionName()))
                .findFirst();

        if (functionCallingOptional.isEmpty()) {
            return Flux.just(new AiMessageResultDTO("调用失败，未找到对应的函数"));
        }

        FunctionCalling functionCalling = functionCallingOptional.get();
        ChatMemory chatMemory = chatMemoryAdvisorProvider.get(chatMessageDTO.getConversationId());
        ChatRequest chatRequest = buildChatRequest(chatMemory, functionCalling, chatMessageDTO.getContent(), jsonAssistantTriple.getRight());
        ChatResponse chatResponse = jsonAssistantTriple.getLeft().chat(chatRequest);
        log.info("json chatResponse: {}", chatResponse.aiMessage().text());

        // 大模型参数解析错误
        if (StrUtil.isBlank(chatResponse.aiMessage().text()) || !JSONUtil.isTypeJSON(chatResponse.aiMessage().text())) {
            return Flux.just(new AiMessageResultDTO("大模型参数解析错误:" + chatResponse.aiMessage().text()));
        }

        // 使用参数调用原有的函数逻辑
        ToolExecutionRequest toolExecutionRequest = ToolExecutionRequest.builder().name(functionCalling.functionName())
                .arguments(chatResponse.aiMessage().text())
                .id(chatMessageDTO.getConversationId())
                .build();

        R<String> resultR = functionCalling.execute(toolExecutionRequest);
        String result = StrUtil.isBlank(resultR.getData()) ? resultR.getMsg() : resultR.getData();
        chatMemory.add(AiMessage.from(chatResponse.aiMessage().text() + StrUtil.LF + result));

        // 返回结果 ,如果业务处理成功的话，增加跳转路由标识
        return Flux.just(new AiMessageResultDTO(result, resultR.isOk() ? functionCalling.routePath() : null));
    }


    /**
     * 构建聊天请求
     *
     * @param chatMemory
     * @param functionCalling 函数调用
     * @param inputText       输入文本
     * @param jsonModel       JSON 模型
     * @return {@link ChatRequest }
     */
    private ChatRequest buildChatRequest(ChatMemory chatMemory, FunctionCalling functionCalling, String inputText, String jsonModel) {
        ChatRequest.Builder builder = ChatRequest.builder();
        Class genericType = functionCalling.getGenericType();
        Field[] fields = ReflectUtil.getFields(genericType);
        JsonObjectSchema jsonObjectSchema = ToolSpecificationsUtils.parametersFrom(fields);
        JsonSchema jsonSchema = JsonSchema.builder()
                .name(functionCalling.functionName())
                .rootElement(jsonObjectSchema)
                .build();

        String metadata = PromptBuilder.render("knowledge-func-metadata.st", Map.of(BaseAiRequest.Fields.messageKey,
                ChatMessageContextHolder.get().getMessageKey(), "systemTime", DateUtil.now()));

        // JSON Schema
        if (ResponseFormatType.JSON_SCHEMA.name().toLowerCase().equals(jsonModel)) {
            chatMemory.add(userMessage(inputText + StrUtil.LF + metadata));
            List<ChatMessage> chatMessages = new ArrayList<>();
            chatMessages.add(systemMessage(PromptBuilder.render("ocr-system-json.st")));
            chatMessages.addAll(chatMemory.messages());
            builder.messages(chatMessages)
                    .responseFormat(ResponseFormat.builder().type(JSON).jsonSchema(jsonSchema).build());
        }

        // user message append json object
        if (ResponseFormatType.JSON_OBJECT.name().toLowerCase().equals(jsonModel)) {
            chatMemory.add(userMessage(PromptBuilder.render("user-json.st",
                    Map.of("userInput", inputText + StrUtil.LF + metadata, "jsonSchema", jsonObjectSchema.toString()))));
            List<ChatMessage> chatMessages = new ArrayList<>();
            chatMessages.add(systemMessage(PromptBuilder.render("ocr-system-json.st")));
            chatMessages.addAll(chatMemory.messages());
            builder.messages(chatMessages);
        }

        return builder.build();
    }

}
